/*---------------------------------------------------------------------
                               头文件
                               erji_app
                             yolo_avsr
-----------------------------------------------------------------------*/

#include "rknn_api.h"
#include "client_t.h"

#include <stdio.h>
#include <errno.h>
#include <cstring>
#include <cstdlib>
#include <unistd.h>
#include "stdlib.h"
#include <pthread.h>
#include <signal.h>
#include <sys/poll.h>
#include <float.h>

#include "rk_defines.h"
#include "rk_debug.h"

#include <vector>
#include "taiic_key.h"
#include <fcntl.h>
#include "taiic_rknn.h"
#include "taiic_sensor.h"
#include "stft.h"

#include <thread>
#include "taiic_opencv.h"
#include "taiic_yolo.h"
#include "taiic_avsr.h"
#include "taiic_mfcc.h"

using namespace std;

#define PICTURE_C_MODEL 49 // 0 +48
#define WIDTH 320
#define HEIGHT 320
#define AICOUNT 26

char aiCmdFinal[256] = {0};

int ai_save_count = 0;

FILE *aifp = NULL;

int ai_save_flag = -1;
int ai_save = 0;

tAiSaveEndCallback *s_ai_save_end = NULL;
/*---------------------------------------------------------------------
                            全局变量/宏定义/结构体/函数声明
-----------------------------------------------------------------------*/

static void sigterm_handler(int sig);

int cnt = 0;

// 判断当前图像是否正在采集
RK_BOOL is_model_run = RK_FALSE;
TAIIC_AI_AO_CTX_S *ctx_aio; // 音频采集、输出相关参数
TAIIC_VI_CTX_S *ctx_vi;     // 视频获取相关配置参数

YOLO_TOOLKIT_MODEL_CTX_S *yolo_ctx; // 模型推理相关
AVSR_TOOLKIT_MODEL_CTX_S *avsr_ctx; // 模型推理相关

TAIIC_MFCC mfccComputer; // audio _ mfcc

RK_U8 vsr_input_data[MODEL_N][MODEL_H][MODEL_W][MODEL_C]; // 唇图输入
RK_U8 frame_data[YOLO_IMG_BATCH][YOLO_IMG_HEIGHT][YOLO_IMG_WIDTH][YOLO_IMG_CHANNEL];
RK_U8 vsr_output_data[1][32][32][1];
RK_U8 avsr_audio_input_data[48][64];

static void power_callback(int power_cap) // 电量查询回调
{
    RK_LOGD("power is %d====\n", power_cap);

}

static void key_callback(TAIIC_KEY_STATE old_state, TAIIC_KEY_NAME key_name, TAIIC_KEY_STATE key_state)
{
    // 判断key_name去实现对应功能,key_name=0表示音频线模式
    RK_LOGE("old_state==%d, key_name==%d, key_state==%d,count=%d\n", old_state, key_name, key_state, cnt++);

    switch (key_name)
    {
    case TAIIC_PTT:

        if (key_state == 0) // 按键按下
        {
            if ((!is_model_run)) // 唇图模式
            {
                is_model_run = RK_TRUE;
                aifp = fopen(aiCmdFinal, "wb");

                system("echo 1 >/sys/class/pwm/pwmchip1/pwm0/enable"); // 开启补光灯
                system("echo b > /sys/devices/virtual/adw/adwdev/adwgpio");
                taiic_vi_get_release_frame_loop(ctx_vi);
                is_model_run = RK_FALSE;
            }
        }

        break;
    }
}

// data 一段音频的指针
// wlen 音频的长度
void ai_frame_data_cback(void *data, RK_U32 wlen)
{
    if (ai_save_flag == 1 && ai_save_count < AICOUNT) // 开始
    {
        // RK_LOGW("=====AI DATA LEN IS %d, count is %d===\n", wlen, ai_save_count);
        if (aifp)
        {
            ai_save_count++;
            fwrite(data, wlen, 1, aifp);
            fflush(aifp);
        }
    }
    else if (ai_save_count == AICOUNT) // 关闭
    {

        fclose(aifp);
        ai_save_flag = -1;
        ai_save_count = 0; // 控制语音启停
        ai_save = 1;       // ai数据保存完成的标志位，go to mfcc

        s_ai_save_end(ai_save);
    }
    else
    {
        return;
    }
}

// data 一帧图像的指针
// wcount 当前图像的顺序
int64_t vi_start_us = 0;
void vi_frame_data_cback(void *data, RK_U32 wcount)
{
    // RK_LOGI("====VI DATA IS %d==\n", wcount);
    if (wcount == 0)
    {
        ai_save_flag = 1; // 语音采集开始
    }
    // ai_save_flag = (wcount == 0) ? 1 : -1;

    // int64_t start_us = getCurrentTimeUs();
    void *tmpData = data;

    // nv12 -> rgb
    taiic_nv12_rgb(tmpData, &frame_data[0][0][0][0], YOLO_IMG_WIDTH, YOLO_IMG_HEIGHT);
    yolo_rknn_toolkit_data_refresh(yolo_ctx, &frame_data[0][0][0][0]);
    int ret = rknn_run(yolo_ctx->context, NULL);
    detect_result_group_t result_group = yolo_rknn_toolkit_result(yolo_ctx);
    RK_LOGD("====rknn _run result num is %d====\n", result_group.count);
    if (result_group.count > 0)
    {
        detect_result_t *det_result = &(result_group.results[1]);
        // RK_LOGD("result is left is %d, top is %d, right is %d, bottom is %d==\n",
        //        det_result->box.left, det_result->box.top,
        //        det_result->box.right, det_result->box.bottom);
        if (det_result->box.left && det_result->box.top && det_result->box.right && det_result->box.bottom)
        {
            taiic_gray_crop_resize(&frame_data[0][0][0][0], &vsr_output_data[0][0][0][0],
                                   YOLO_IMG_WIDTH, YOLO_IMG_HEIGHT,
                                   det_result->box.left, det_result->box.top,
                                   det_result->box.right, det_result->box.bottom,
                                   32, 32);
        }
        else
        {
            taiic_gray_resize(&frame_data[0][0][0][0], &vsr_output_data[0][0][0][0],
                              YOLO_IMG_WIDTH, YOLO_IMG_HEIGHT,
                              32, 32);
        }
    }
    else
    {
        taiic_gray_resize(&frame_data[0][0][0][0], &vsr_output_data[0][0][0][0],
                          YOLO_IMG_WIDTH, YOLO_IMG_HEIGHT,
                          32, 32);
    }
    // free(&frame_data[0][0][0][0]);

    RK_U8 *pTmpData = &vsr_output_data[0][0][0][0];
    int cnow = wcount / (RESHAPE_H * RESHAPE_W);               // 通道数
    int hnow = (wcount % (RESHAPE_H * RESHAPE_W)) / RESHAPE_W; // 高
    int wnow = (wcount % (RESHAPE_H * RESHAPE_W)) / RESHAPE_H; // 宽
    for (int i = 0; i < RESIZE_PIC_H; i++)
    {
        for (int j = 0; j < RESIZE_PIC_W; j++)
        {
            vsr_input_data[0][RESIZE_PIC_H * hnow + i][RESIZE_PIC_W * wnow + j][cnow] = *pTmpData;
            pTmpData += 1;
        }
    }

    // int64_t elapse_us = getCurrentTimeUs() - start_us;
    // RK_LOGD("Elapse Time = %.2fms, FPS = %.2f\n", elapse_us / 1000.f, 1000.f * 1000.f / elapse_us);

    int64_t vi_frame_elapse_us = getCurrentTimeUs() - vi_start_us;
    vi_start_us = getCurrentTimeUs();
    RK_LOGD("Vi Frame %d Elapse Time = %.2fms, FPS = %.2f\n", wcount, vi_frame_elapse_us / 1000.f, 1000.f * 1000.f / vi_frame_elapse_us);
}

static void ai_save_end(int ai)
{
    ai_save = 0;
    RK_LOGD("=====save ai end==\n");
    ////end save ai -> go to mfcc
    ////ai_data in /tmp/test.pcm 16k 16bits 1channel
    std::ifstream wavFp;
    std::ofstream mfcFp;

    // Check if input is readable
    wavFp.open("/tmp/test.pcm");
    if (!wavFp.is_open())
    {
        std::cerr << "Unable to open input file: " << std::endl;
    }
    mfccComputer.process(wavFp, mfcFp, &avsr_audio_input_data[0][0]);

    unsigned char *avsr_input_data[avsr_ctx->io_num.n_input];
    avsr_input_data[0] = new unsigned char[avsr_ctx->input_attrs[0].size]; // 设置video 输入大小
    avsr_input_data[0] = &vsr_input_data[0][0][0][0];

    avsr_input_data[1] = new unsigned char[avsr_ctx->input_attrs[1].size]; // 设置audio 输入大小
    avsr_input_data[1] = &avsr_audio_input_data[0][0];

    avsr_rknn_toolkit_data_refresh(avsr_ctx, avsr_input_data);
    int ret = rknn_run(avsr_ctx->context, NULL);
    MODEL_RESULT_S rknn_result = avsr_rknn_toolkit_result(avsr_ctx);
    RK_LOGD("===result label is %d, prob is %f====\n", rknn_result.label, rknn_result.prob);

    // for (int i = 0; i < avsr_ctx->io_num.n_input; i++)
    // {
    //     free(avsr_input_data[i]);
    // }
}

void ai_save_end_callback(tAiSaveEndCallback ai_save_end)
{
    if (ai_save_end != NULL)
    {
        s_ai_save_end = ai_save_end;
    }
}

/*----------------------------------------------------------------
                            主函数main
-----------------------------------------------------------------*/

int main(int argc, char **argv)
{
    system("mic0_taiic.sh");
    system("echo f > /sys/devices/virtual/adw/adwdev/adwgpio"); // mic_mic0

    // initial aio
    ctx_aio = reinterpret_cast<TAIIC_AI_AO_CTX_S *>(malloc(sizeof(TAIIC_AI_AO_CTX_S)));
    memset(ctx_aio, 0, sizeof(TAIIC_AI_AO_CTX_S));

    // initial vi
    ctx_vi = reinterpret_cast<TAIIC_VI_CTX_S *>(malloc(sizeof(TAIIC_VI_CTX_S))); // 分配内存空间
    memset(ctx_vi, 0, sizeof(TAIIC_VI_CTX_S));                                   // 统一设置为0

    // initial yolo model
    yolo_ctx = reinterpret_cast<YOLO_TOOLKIT_MODEL_CTX_S *>(malloc(sizeof(YOLO_TOOLKIT_MODEL_CTX_S))); // 分配内存空间
    memset(yolo_ctx, 0, sizeof(YOLO_TOOLKIT_MODEL_CTX_S));

    // initial avsr model
    avsr_ctx = reinterpret_cast<AVSR_TOOLKIT_MODEL_CTX_S *>(malloc(sizeof(AVSR_TOOLKIT_MODEL_CTX_S))); // 分配内存空间
    memset(avsr_ctx, 0, sizeof(AVSR_TOOLKIT_MODEL_CTX_S));

    snprintf(aiCmdFinal, sizeof(aiCmdFinal), "/tmp/test.pcm");

    // config ai
    taiic_ai_ao_media_config_init(ctx_aio);
    ctx_aio->aio_trans.srcData = reinterpret_cast<RK_U8 *>(calloc(ctx_aio->ao_size, sizeof(RK_U8)));
    ctx_aio->s32SetTrackMode = 8; //

    // config vi
    taiic_vi_media_config_init(ctx_vi);
    ctx_vi->width = WIDTH;
    ctx_vi->height = HEIGHT;
    ctx_vi->loopCountSet = PICTURE_C_MODEL;

    // config yolo model
    yolo_ctx->modelPath = "/oem/usr/model/yolo.rknn";
    RK_LOGI("==This is line %d in files==\n", __LINE__);
    yolo_rknn_toolkit_config_init(yolo_ctx);
    yolo_rknn_toolkit_io_init(yolo_ctx);

    avsr_ctx->modelPath = "/oem/usr/model/avsr.rknn"; // 模型存放路径
    avsr_rknn_toolkit_config_init(avsr_ctx);
    avsr_rknn_toolkit_io_init(avsr_ctx);

    RK_LOGI("------------params set------------\n");

    // ai _save end_callback
    ai_save_end_callback(ai_save_end);
    // 按键检测
    taiic_all_key_registers(ctx_aio, key_callback);
    // 电量查询检测回调
    taiic_batt_power_registers(power_callback);

    RK_LOGI("====go to callback===\n");
    taiic_ai_frame_data_callback(ai_frame_data_cback); // 注册ai数据回调事件

    taiic_vi_frame_data_callback(vi_frame_data_cback); // 注册vi数据回调事件

    // setup the SIGINT to ctrl+c and handing de-init process
    signal(SIGINT, sigterm_handler);
    RK_MPI_SYS_Init();

    taiic_aio_data_capture(ctx_aio);

__FAILED_AIO:
    if (ctx_aio)
    {
        free(ctx_aio);
        ctx_aio = RK_NULL;
    }
__FAILED_VI:
    if (ctx_vi)
    {
        free(ctx_vi);
        ctx_vi = RK_NULL;
    }
__FAILED_YOLO:
    if(yolo_ctx)
    {
        free(yolo_ctx);
        yolo_ctx = RK_NULL;
    }
__FAILED_AVSR:
    if(avsr_ctx)
    {
        free(avsr_ctx);
        avsr_ctx = RK_NULL;
    }

    RK_MPI_SYS_Exit();
    return 0;
}
/*---------------------------------------------------------------
                          函数实现
-----------------------------------------------------------------*/

static void sigterm_handler(int sig)
{
    RK_LOGI("Catched SIGINT %d\n", sig);
    ctx_aio->gAiExit = RK_TRUE;
}
